#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <chrono>

#pragma comment(lib, "ws2_32.lib")

#define BUFFER_SIZE 1024 * 1024 // 1MB缓冲区，可根据实际情况调整
#define SERVER_IP "127.0.0.1"
#define PORT 8888

// TCP发送文件函数
void tcpSendFile(const std::string& filePath) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
        return;
    }

    SOCKET sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock == INVALID_SOCKET) {
        std::cerr << "Socket creation failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return;
    }

    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(PORT);
    InetPton(AF_INET, SERVER_IP, &serverAddr.sin_addr);

    if (connect(sock, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
        std::cerr << "Connect failed: " << WSAGetLastError() << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    std::ifstream file(filePath, std::ios::binary);
    if (!file) {
        std::cerr << "Failed to open file: " << filePath << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    auto start = std::chrono::high_resolution_clock::now();
    char buffer[BUFFER_SIZE];
    while (file.read(buffer, sizeof(buffer))) {
        if (send(sock, buffer, sizeof(buffer), 0) == SOCKET_ERROR) {
            std::cerr << "Send failed: " << WSAGetLastError() << std::endl;
            break;
        }
    }
    // 发送剩余数据
    std::streamsize remaining = file.gcount();
    if (remaining > 0) {
        file.read(buffer, remaining);
        if (send(sock, buffer, remaining, 0) == SOCKET_ERROR) {
            std::cerr << "Send failed: " << WSAGetLastError() << std::endl;
        }
    }

    file.close();
    closesocket(sock);
    WSACleanup();

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "TCP传输时间: " << duration.count() << " 秒" << std::endl;
}

// TCP接收文件函数
void tcpReceiveFile(const std::string& savePath) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
        return;
    }

    SOCKET listenSock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listenSock == INVALID_SOCKET) {
        std::cerr << "Socket creation failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return;
    }

    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(PORT);
    serverAddr.sin_addr.s_addr = INADDR_ANY;

    if (bind(listenSock, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
        std::cerr << "Bind failed: " << WSAGetLastError() << std::endl;
        closesocket(listenSock);
        WSACleanup();
        return;
    }

    if (listen(listenSock, 5) == SOCKET_ERROR) {
        std::cerr << "Listen failed: " << WSAGetLastError() << std::endl;
        closesocket(listenSock);
        WSACleanup();
        return;
    }

    sockaddr_in clientAddr;
    int clientAddrLen = sizeof(clientAddr);
    SOCKET clientSock = accept(listenSock, (sockaddr*)&clientAddr, &clientAddrLen);
    if (clientSock == INVALID_SOCKET) {
        std::cerr << "Accept failed: " << WSAGetLastError() << std::endl;
        closesocket(listenSock);
        WSACleanup();
        return;
    }

    std::ofstream file(savePath, std::ios::binary);
    if (!file) {
        std::cerr << "Failed to open file for writing: " << savePath << std::endl;
        closesocket(clientSock);
        closesocket(listenSock);
        WSACleanup();
        return;
    }

    auto start = std::chrono::high_resolution_clock::now();
    char buffer[BUFFER_SIZE];
    int bytesReceived;
    while ((bytesReceived = recv(clientSock, buffer, sizeof(buffer), 0)) > 0) {
        file.write(buffer, bytesReceived);
    }

    file.close();
    closesocket(clientSock);
    closesocket(listenSock);
    WSACleanup();

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "TCP接收时间: " << duration.count() << " 秒" << std::endl;
}

// UDP发送文件函数
void udpSendFile(const std::string& filePath) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
        return;
    }

    SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET) {
        std::cerr << "Socket creation failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return;
    }

    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(PORT);
    InetPton(AF_INET, SERVER_IP, &serverAddr.sin_addr);

    std::ifstream file(filePath, std::ios::binary);
    if (!file) {
        std::cerr << "Failed to open file: " << filePath << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    auto start = std::chrono::high_resolution_clock::now();
    char buffer[BUFFER_SIZE];
    while (file.read(buffer, sizeof(buffer))) {
        if (sendto(sock, buffer, sizeof(buffer), 0, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
            std::cerr << "Sendto failed: " << WSAGetLastError() << std::endl;
            break;
        }
    }
    // 发送剩余数据
    std::streamsize remaining = file.gcount();
    if (remaining > 0) {
        file.read(buffer, remaining);
        if (sendto(sock, buffer, remaining, 0, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
            std::cerr << "Sendto failed: " << WSAGetLastError() << std::endl;
        }
    }

    file.close();
    closesocket(sock);
    WSACleanup();

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "UDP传输时间: " << duration.count() << " 秒" << std::endl;
}

// UDP接收文件函数
void udpReceiveFile(const std::string& savePath) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
        return;
    }

    SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET) {
        std::cerr << "Socket creation failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return;
    }

    sockaddr_in localAddr;
    localAddr.sin_family = AF_INET;
    localAddr.sin_port = htons(PORT);
    localAddr.sin_addr.s_addr = INADDR_ANY;

    if (bind(sock, (sockaddr*)&localAddr, sizeof(localAddr)) == SOCKET_ERROR) {
        std::cerr << "Bind failed: " << WSAGetLastError() << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    std::ofstream file(savePath, std::ios::binary);
    if (!file) {
        std::cerr << "Failed to open file for writing: " << savePath << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    auto start = std::chrono::high_resolution_clock::now();
    char buffer[BUFFER_SIZE];
    sockaddr_in senderAddr;
    int senderAddrLen = sizeof(senderAddr);
    int bytesReceived;
    while ((bytesReceived = recvfrom(sock, buffer, sizeof(buffer), 0, (sockaddr*)&senderAddr, &senderAddrLen)) > 0) {
        file.write(buffer, bytesReceived);
    }

    file.close();
    closesocket(sock);
    WSACleanup();

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "UDP接收时间: " << duration.count() << " 秒" << std::endl;
}

int main() {
    // 示例：TCP发送文件
    tcpSendFile("large_file.exe");
    // 示例：TCP接收文件
    tcpReceiveFile("received_file_tcp.exe");

    // 示例：UDP发送文件
    udpSendFile("large_file.exe");
    // 示例：UDP接收文件
    udpReceiveFile("received_file_udp.exe");

    return 0;
}    